using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Diagnostics;
using System.Linq;

namespace EvoBoost.Core
{
	internal class TreeNode<T>
	{
		private List<double> splitParameters;

		private ReadOnlyCollection<double> splitParametersReadOnly;

		private List<double> previousOutcomes;

		private List<Sample<T>> samples;

		private bool isDirty;

		private double thresholdOrOutcome;

		private double totalLoss;

		private TreeNode()
		{
			this.Type = TreeNodeType.Terminal;
			this.TerminalsInSubtree = 1;
			this.isDirty = true;
		}

		public static TreeNode<T> CreateRoot(
			ICollection<double> previousOutcomes,
			ICollection<Sample<T>> samples)
		{
			TreeNode<T> root = new TreeNode<T>();
			root.previousOutcomes = new List<double>(previousOutcomes);
			root.samples = new List<Sample<T>>(samples);
			return root;
		}

		public int TerminalsInSubtree { get; private set; }

		public TreeNode<T> Parent { get; private set; }

		public TreeNode<T> LeftChild { get; private set; }

		public TreeNode<T> RightChild { get; private set; }

		public TreeNodeType Type { get; private set; }

		public double ThresholdOrOutcome
		{
			get
			{
				Debug.Assert(!this.isDirty);
				return this.thresholdOrOutcome;
			}
		}

		/// <summary>
		/// Gets total loss sum in subtree.
		/// </summary>
		public double TotalLoss
		{
			get
			{
				Debug.Assert(!this.isDirty);
				return this.totalLoss;
			}
		}

		public ReadOnlyCollection<double> SplitParameters
		{
			get { return this.splitParametersReadOnly; }
		}

		private void UpdateTerminalsInSubtree()
		{
			if (this.Type == TreeNodeType.Terminal)
				this.TerminalsInSubtree = 1;
			else if (this.Type == TreeNodeType.Splitter)
				this.TerminalsInSubtree = this.LeftChild.TerminalsInSubtree + this.RightChild.TerminalsInSubtree;
			else
				Debug.Fail("We should never get there.");

			if (this.Parent != null)
				this.Parent.UpdateTerminalsInSubtree();
		}

		public void RefreshParametersInSubtree(ILossFunction lossFunction, SplitBase<T> split, int minSamplesInNode)
		{
			if (this.isDirty)
			{
				// Prepare list of outcomes
				List<double> outcomes = new List<double>(samples.Count);
				foreach (Sample<T> sample in this.samples)
					outcomes.Add(sample.Outcome);

				if (this.Type == TreeNodeType.Splitter)
				{
					// Calculate split values
					List<double> splitValues = new List<double>(this.samples.Count);
					foreach (Sample<T> sample in this.samples)
						splitValues.Add(split.Calculate(sample.Properties, this.SplitParameters));
					
					// Calculate threshold for splitter
					this.thresholdOrOutcome = lossFunction.FindOptimalSplit(this.previousOutcomes, outcomes, splitValues);

					// Split samples
					List<Sample<T>> samplesForLeftNode = new List<Sample<T>>(this.samples.Count);
					List<double> prevOutcomesForLeftNode = new List<double>(this.samples.Count);
					List<Sample<T>> samplesForRightNode = new List<Sample<T>>(this.samples.Count);
					List<double> prevOutcomesForRightNode = new List<double>(this.samples.Count);
					for (int i = 0; i < this.samples.Count; ++i)
					{
						if (splitValues[i] <= this.thresholdOrOutcome)
						{
							samplesForLeftNode.Add(this.samples[i]);
							prevOutcomesForLeftNode.Add(this.previousOutcomes[i]);
						}
						else
						{
							samplesForRightNode.Add(this.samples[i]);
							prevOutcomesForRightNode.Add(this.previousOutcomes[i]);
						}
					}

					// Check how much samples we have in node
					if (samplesForLeftNode.Count < minSamplesInNode || samplesForRightNode.Count < minSamplesInNode)
					{
						// Collapse node cause we should not split it (too few samples available)
						this.Collapse();
					}
					else
					{
						this.LeftChild.samples = samplesForLeftNode;
						this.LeftChild.previousOutcomes = prevOutcomesForLeftNode;
						this.LeftChild.isDirty = true;

						this.RightChild.samples = samplesForRightNode;
						this.RightChild.previousOutcomes = prevOutcomesForRightNode;
						this.RightChild.isDirty = true;
					}
				}

				// Node was or become a terminal?
				if (this.Type == TreeNodeType.Terminal)
				{
					this.thresholdOrOutcome = lossFunction.Minimize(this.previousOutcomes, outcomes);

					Debug.Assert(this.samples.Count == this.previousOutcomes.Count);
					this.totalLoss = 0;
					for (int i = 0; i < this.samples.Count; ++i)
					{
						this.totalLoss += lossFunction.Calculate(
							this.samples[i].Outcome, this.previousOutcomes[i] + this.thresholdOrOutcome);
					}
				}

				// Anyway, it is not dirty now
				this.isDirty = false;
			}

			if (this.Type == TreeNodeType.Splitter)
			{
				this.LeftChild.RefreshParametersInSubtree(lossFunction, split, minSamplesInNode);
				this.RightChild.RefreshParametersInSubtree(lossFunction, split, minSamplesInNode);

				this.totalLoss = this.LeftChild.TotalLoss + this.RightChild.TotalLoss;
			}
		}

		public static void SwapSubtrees(TreeNode<T> node1, TreeNode<T> node2)
		{
			Debug.Assert(node1 != null && node1.Parent != null);
			Debug.Assert(node2 != null && node2.Parent != null);

			TreeNode<T> node1Parent = node1.Parent;
			TreeNode<T> node2Parent = node2.Parent;
			bool is1Left = node1Parent.LeftChild == node1;
			bool is2Left = node2Parent.LeftChild == node2;

			// Swap parents
			node1.Parent = node2Parent;
			node2.Parent = node1Parent;

			// Connect 1st node to the new parent
			if (is2Left)
				node1.Parent.LeftChild = node1;
			else
				node1.Parent.RightChild = node1;

			// Connect 2nd node to the new parent
			if (is1Left)
				node2.Parent.LeftChild = node2;
			else
				node2.Parent.RightChild = node2;

			// Update information about terminal node count
			node1.Parent.UpdateTerminalsInSubtree();
			node2.Parent.UpdateTerminalsInSubtree();

			// Update dirty state
			node1.Parent.isDirty = true;
			node2.Parent.isDirty = true;
		}

		public void Split(IEnumerable<double> parameters)
		{
			Debug.Assert(this.Type == TreeNodeType.Terminal && this.TerminalsInSubtree == 1);

			this.Type = TreeNodeType.Splitter;
			this.splitParameters = new List<double>(parameters);
			this.splitParametersReadOnly = new ReadOnlyCollection<double>(this.splitParameters);

			TreeNode<T> leftChild = new TreeNode<T>();
			this.LeftChild = leftChild;
			this.LeftChild.Parent = this;
			TreeNode<T> rightChild = new TreeNode<T>();
			this.RightChild = rightChild;
			this.RightChild.Parent = this;

			this.UpdateTerminalsInSubtree();

			Debug.Assert(this.TerminalsInSubtree == 2);

			this.isDirty = true;
		}

		public void Adjust(IEnumerable<double> parameters)
		{
			Debug.Assert(this.Type == TreeNodeType.Splitter);

			this.splitParameters.Clear();
			this.splitParameters.AddRange(parameters);

			this.isDirty = true;
		}

		public void Collapse()
		{
			Debug.Assert(this.Type == TreeNodeType.Splitter);

			this.Type = TreeNodeType.Terminal;
			this.splitParameters = null;
			this.splitParametersReadOnly = null;
			this.LeftChild.Parent = null;
			this.LeftChild = null;
			this.RightChild.Parent = null;
			this.RightChild = null;

			this.UpdateTerminalsInSubtree();

			this.isDirty = true;
		}

		public void ExtractSubtree(
			ICollection<TreeNode<T>> allNodes,
			ICollection<TreeNode<T>> terminalNodes,
			ICollection<TreeNode<T>> splitterNodes)
		{
			if (allNodes != null) allNodes.Add(this);

			if (this.Type == TreeNodeType.Terminal)
			{
				if (terminalNodes != null) terminalNodes.Add(this);
			}
			else if (this.Type == TreeNodeType.Splitter)
			{
				if (splitterNodes != null) splitterNodes.Add(this);

				this.LeftChild.ExtractSubtree(allNodes, terminalNodes, splitterNodes);
				this.RightChild.ExtractSubtree(allNodes, terminalNodes, splitterNodes);
			}
		}

		public TreeNode<T> CloneSubtree()
		{
			return CloneSubtreeDfs(this);
		}

		private static TreeNode<T> CloneSubtreeDfs(TreeNode<T> nodeToCopy)
		{
			// Create node clone
			TreeNode<T> result = new TreeNode<T>
			{
				Parent = null,
				Type = nodeToCopy.Type,
				TerminalsInSubtree = nodeToCopy.TerminalsInSubtree,
				thresholdOrOutcome = nodeToCopy.thresholdOrOutcome,
				totalLoss = nodeToCopy.totalLoss,
				isDirty = nodeToCopy.isDirty,
				previousOutcomes = new List<double>(nodeToCopy.previousOutcomes),
				samples = new List<Sample<T>>(nodeToCopy.samples),
			};

			// Clone subtrees
			if (nodeToCopy.Type == TreeNodeType.Splitter)
			{
				result.LeftChild = CloneSubtreeDfs(nodeToCopy.LeftChild);
				result.LeftChild.Parent = result;
				result.RightChild = CloneSubtreeDfs(nodeToCopy.RightChild);
				result.RightChild.Parent = result;
				result.splitParameters = new List<double>(nodeToCopy.splitParameters);
				result.splitParametersReadOnly = new ReadOnlyCollection<double>(result.splitParameters);
			}

			return result;
		}
	}
}